import networkx as nx
import numpy as np
from neo4j import GraphDatabase

SCENARIO_CONFIG = {
    '核心供应商': {
        'name': "1-核心供应商依赖",
        'desc': "关键零部件高度依赖单一供应商，供应中断时产生多米诺效应",
        'params': {'alpha': 0.85, 'beta': 0.35, 'conn_density': (5, 8),
                   'industry_impact': {'制造': {'base': 4500, 'sigma': 0.15}}}
    },
    '技术封锁': {
        'name': "2-技术封锁",
        'desc': "关键技术专利受限导致生产链关键环节断裂",
        'params': {'alpha': 0.9, 'beta': 0.25, 'conn_density': (3, 5),
                   'industry_impact': {'制造': {'base': 2000, 'sigma': 0.4}}}
    },
    '政治冲突': {
        'name': "3-政治冲突",
        'desc': "地缘政治冲突导致跨境物流通道受阻",
        'params': {'alpha': 0.75, 'beta': 0.45, 'conn_density': (2, 4),
                   'industry_impact': {'制造': {'base': 3000, 'sigma': 0.5}}}
    },
    '区域热点': {
        'name': "4-局部冲突",
        'desc': "地区武装冲突导致基础设施损毁和运输中断",
        'params': {'alpha': 0.95, 'beta': 0.3, 'conn_density': (1, 3),
                   'industry_impact': {'制造': {'base': 5000, 'sigma': 0.6}}}
    },
    '自然灾害': {
        'name': "5-自然灾害",
        'desc': "极端天气事件造成区域供应链节点瘫痪",
        'params': {'alpha': 0.8, 'beta': 0.4, 'conn_density': (4, 6),
                   'industry_impact': {'制造': {'base': 1500, 'sigma': 0.7}}}
    },
    '生产经营': {
        'name': "6-经营风险",
        'desc': "安全事故引发上下游企业连锁反应",
        'params': {'alpha': 0.7, 'beta': 0.5, 'conn_density': (6, 8),
                   'industry_impact': {'制造': {'base': 2500, 'sigma': 0.3}}}
    },
    '质量安全': {
        'name': "7-质量风险",
        'desc': "产品缺陷引发全产业链召回和信任危机",
        'params': {'alpha': 0.85, 'beta': 0.35, 'conn_density': (5, 7),
                   'industry_impact': {'制造': {'base': 1800, 'sigma': 0.4}}}
    }
}

RISK_ANALYSIS = {
    '核心供应商': "核心供应商依赖导致产业链关键节点失效",
    '技术封锁': "技术瓶颈引发创新链断裂风险",
    '政治冲突': "地缘政治冲击跨境供应链",
    '区域热点': "区域冲突造成供应链物理中断",
    '自然灾害': "自然灾害破坏区域物流网络",
    '生产经营': "安全事故风险传导至生产网络",
    '质量安全': "质量危机引发信任链崩溃"
}


class AdvancedSupplyChainModel:
    def __init__(self,
                 neo4j_uri="bolt://localhost:7687",
                 neo4j_user="neo4j",
                 neo4j_password="neo4j",
                 use_neo4j=True):
        """
        初始化模型（不立即加载数据）
        """
        self.G = nx.DiGraph()
        self.use_neo4j = use_neo4j
        self.project_id = None
        self.enterprises = []
        self.driver = None

        if self.use_neo4j:
            self.driver = GraphDatabase.driver(neo4j_uri, auth=(neo4j_user, neo4j_password))

    def close(self):
        if self.driver:
            self.driver.close()

    def _fetch_enterprises_from_neo4j(self, project_id):
        query = """
        MATCH (n) 
        WHERE n.name IS NOT NULL AND n.projectId = $projectId
        RETURN n.name AS name
        """
        with self.driver.session() as session:
            result = session.run(query, projectId=project_id)
            return [record["name"] for record in result]

    def _init_network_from_neo4j(self, project_id):
        self.G.clear()
        self.G.add_nodes_from(self.enterprises)

        query = """
        MATCH (s)-[r{name:'供应'}]->(m)         
        WHERE s.projectId = $projectId AND m.projectId = $projectId
        RETURN s.name AS supplier, m.name AS main, r.width AS weight
        """

        with self.driver.session() as session:
            result = session.run(query, projectId=project_id)
            for record in result:
                supplier = record["supplier"]
                main = record["main"]
                weight = record.get("weight", np.random.lognormal(mean=3, sigma=0.8))
                if supplier in self.enterprises and main in self.enterprises:
                    self.G.add_edge(supplier, main, weight=weight)

    def _load_project_data_if_needed(self, project_id):
        self.project_id = project_id
        self.enterprises = self._fetch_enterprises_from_neo4j(project_id)
        self._init_network_from_neo4j(project_id)
        print(f"成功加载项目 {project_id} 的供应链数据")
        print(f"节点数量：{len(self.enterprises)}，边数量：{self.G.number_of_edges()}")
        print(f"节点样本：{self.enterprises[:10]}")



    def simulate_cascade(self, scenario_id, start_enterprise, project_id):
        """
        模拟供应链级联失效过程
        """
        print(f"开始模拟 {scenario_id} 情景")
        print(f"起始企业：{start_enterprise}")
        print(f"项目编号：{project_id}")
        if scenario_id not in SCENARIO_CONFIG:
            raise ValueError(f"无效情景编号：{scenario_id}")

        if self.use_neo4j:
            self._load_project_data_if_needed(project_id)

        if start_enterprise not in self.enterprises:
            raise ValueError(f"企业 '{start_enterprise}' 不存在于供应链中")

        params = SCENARIO_CONFIG[scenario_id]['params']
        failed = {start_enterprise}
        queue = [start_enterprise]

        while queue:
            current = queue.pop(0)
            for successor in self.G.successors(current):
                if successor in failed:
                    continue
                edge_data = self.G.get_edge_data(current, successor)
                impact = params['alpha'] * edge_data['weight']
                total_inflow = sum(self.G[pre][successor]['weight'] for pre in self.G.predecessors(successor))
                threshold = params['beta'] * total_inflow
                if impact > threshold:
                    failed.add(successor)
                    queue.append(successor)

        return {
            'scenario': SCENARIO_CONFIG[scenario_id]['name'],
            'description': SCENARIO_CONFIG[scenario_id]['desc'],
            'start_enterprise': start_enterprise,
            'failed_enterprises': sorted(failed),
            'risk_analysis': RISK_ANALYSIS[scenario_id]
        }
